from CNN import ConvNet
import torch
import torchvision.transforms as transform
import matplotlib.pyplot as plt
from PIL import Image
from torch.autograd import Variable
from tool import *
from DATASET import datasetting

cnn=ConvNet()
cnn=torch.load('checkpoint/minst_conv_checkpoint')

#mnist是黑底白字
test_img='data/bdhz.png'




#########将黑白反转——白底黑字-》黑底白字
img=Image.open(test_img)
img=convertL(img)
img = img.resize((28, 28))

# 暂存像素值的一维数组
arr = []

for i in range(28):
    for j in range(28):
        # mnist 里的颜色是0代表白色（背景），1.0代表黑色
        pixel = 1.0 - float(img.getpixel((j, i)))/255.0
        # pixel = 255.0 - float(img.getpixel((j, i))) # 如果是0-255的颜色值
        arr.append(pixel)

arr1 = np.array(arr).reshape((1,28,28))
show_arrimg(arr1)

arr1=torch.from_numpy(arr1)
arr1=arr1.unsqueeze(0)

dtpe=torch.cuda.FloatTensor
arr1=arr1.type(dtpe)

output=cnn(arr1)
print(output)
pred=torch.max(output.data,1)[1]
print(pred.item())


#直接进行测试
# PILimg=Image.open(test_img)
# img=convertL(PILimg)
# loader=transform.Compose([transform.Resize(28),transform.CenterCrop((28,28)),transform.ToTensor()])
# img=loader(img)
# show_tensorimg(img)
#
# img=img.unsqueeze(0)
# print (img.size(),img.shape,img.type)
# img=img.cuda()
#
# output=cnn(img)
# print(output)
# pred=torch.max(output.data,1)[1]
# print(pred.item())

#
#
######minst 数字3 ##############
# _,_,test=datasetting()
# idx=93
# mt=test.dataset[idx][0].numpy()
# plt.imshow(mt[0,...])
# plt.show()
# print(test.dataset[idx][1])
#
# mt=torch.from_numpy(mt)
# mt=Variable(mt)
#
#
# mt=mt.cuda()
# mt=mt.unsqueeze(0)
# print(mt.size())
#
# output=cnn(mt)
# print(output)
# pred=torch.max(output.data,1)[1]
# print(pred.item())

